package test.datasource;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;

/**
 * @author mh.z
 */
public class UnitTestDataSource extends DataSourceWrapper {
    private UnitTestConnection currentConnection;

    public static final String UNIT_TEST_DATA_SOURCE_NAME = "unitTestDataSource";

    public UnitTestDataSource(DataSource targetDataSource) {
        super(targetDataSource);
    }

    public UnitTestConnection getCurrentConnection() {
        return currentConnection;
    }

    @Override
    public Connection getConnection() throws SQLException {
        if (currentConnection != null) {
            return currentConnection;
        }

        Connection newConnection = super.getConnection();
        currentConnection = new UnitTestConnection(newConnection);

        return currentConnection;
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        if (currentConnection != null) {
            return currentConnection;
        }

        Connection newConnection = super.getConnection(username, password);
        currentConnection = new UnitTestConnection(newConnection);

        return currentConnection;
    }

    /**
     * 关闭数据源
     */
    public void close() {
        if (currentConnection != null) {
            currentConnection.forceClose();
            currentConnection = null;
        }
    }

    /**
     * 测试前置处理
     */
    public void setUp() {
        if (currentConnection != null) {
            currentConnection.forceRollback();
        }
    }

    /**
     * 测试后置处理
     */
    public void tearDown() {
        if (currentConnection != null) {
            currentConnection.forceRollback();
        }
    }

}
